/*
 * EquilibriumProblem.cpp
 *
 *  Created on: 27 Jul 2011
 *      Author: Allan
 */

#include "EquilibriumProblem.h"

// GeoReact includes
#include "EquilibriumSolver.h"

EquilibriumProblem::EquilibriumProblem()
{}

EquilibriumProblem::EquilibriumProblem(const Multiphase& multiphase, const ReactionSystem& eReactions)
{
	Initialiaze(multiphase, eReactions);
}

EquilibriumProblem::~EquilibriumProblem()
{}

void EquilibriumProblem::Initialiaze(const Multiphase& multiphase, const ReactionSystem& eReactions)
{
	// Set data members
	this->multiphase = multiphase;
	this->eReactions = eReactions;
	
	// Set the number of primary (Nj) and equilibrium (Ne) species
	Ni = multiphase.GetNumSpecies();
	Ne = eReactions.GetNumReactions();
	Nj = Ni - Ne;
	
	// Set the stoichiometric matrix of the system of reactions
	uei = eReactions.AssembleStoichiometricMatrix(multiphase.GetSpecies());
	
	// Allocate all the auxiliary vector/matrix variables with zero entries
	Ke = VectorXd::Zero(Ne);
	h  = VectorXd::Zero(Nj);
	
	alpha = MatrixXd::Zero(Nj, Ni);
	
	hfunctions.resize(Nj);
}

void EquilibriumProblem::SetEquilibriumConditions(const vector<EquilibriumCondition>& eConditions)
{
	// Loop over all the equilibrium conditions
	for(unsigned j = 0; j < Nj; ++j)
	{
		// Set the j-th row of the alpha matrix
		alpha.row(j) = eConditions[j].alphaRow;
		
		// Set the j-th row of the equilibrium condition functions
		hfunctions[j] = eConditions[j].h;
	}
}

void EquilibriumProblem::SetParameters(double T_, double P_)
{
	// Set the parameters temperature (T) and pressure (P)
	T = T_;	
	P = P_;
	
	// Set the equilibrium constants of the equilibrium reactions (Ke)
	Ke = eReactions.EquilibriumConstants(T, P);
}

void EquilibriumProblem::Function(const VectorXd& n, VectorXd& F)
{
	// Check if all the entries in (njne) are non-negative
	if(n.minCoeff() > 0)
	{
		// Update the activity vector (a) with temperature (T), pressure (P) and the mole vector (n)
		a = multiphase.Activities(T, P, n);
		
		// Update the vector (h), the right-hand side values of the equilibrium conditions
		for(unsigned j = 0; j < Nj; ++j) 
			h[j] = hfunctions[j](T, P, n, a);
	}
	
	// Calculate the logarithm of the vector of activities (a) and equilibrium constants (Ke)
	const VectorXd ln_a  =  a.array().log();
	const VectorXd ln_Ke = Ke.array().log();
	
	// Compute the residual vector Fe
	F.segment(0, Ne) = uei * ln_a - ln_Ke;
	
	// Compute the residual vector Fj
	F.segment(Ne, Nj) = alpha * n - h;
}

void EquilibriumProblem::Jacobian(const VectorXd& n, MatrixXd& J)
{
	// Compute the block Jacobian matrix Jei
	J.block(0, 0, Ne, Ni) = uei * n.asDiagonal().inverse();
	
	// Compute the block Jacobian matrix Jji
	J.block(Ne, 0, Nj, Ni) = alpha;
}
